package cn.sytton.taffeauth.config;

import org.springframework.security.core.session.SessionInformation;
import org.springframework.security.crypto.keygen.StringKeyGenerator;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.jose.jws.JwsAlgorithm;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.oauth2.jwt.*;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.UUID;

public class UUIDOAuth2IdTokenGenerator implements OAuth2TokenGenerator<Jwt> {
    private final StringKeyGenerator accessTokenGenerator = new UUIDKeyGenerator();

    private final JwtEncoder jwtEncoder;

    public UUIDOAuth2IdTokenGenerator(JwtEncoder jwtEncoder) {
        Assert.notNull(jwtEncoder, "jwtEncoder cannot be null");
        this.jwtEncoder = jwtEncoder;
    }

    @Override
    public Jwt generate(OAuth2TokenContext context) {
        if (context.getTokenType() != null && "id_token".equals(context.getTokenType().getValue())) {
            String issuer = null;
            if (context.getAuthorizationServerContext() != null) {
                issuer = context.getAuthorizationServerContext().getIssuer();
            }

            RegisteredClient registeredClient = context.getRegisteredClient();
            Instant issuedAt = Instant.now();
            JwsAlgorithm jwsAlgorithm = SignatureAlgorithm.RS256;
            Instant expiresAt;

            expiresAt = issuedAt.plus(30L, ChronoUnit.MINUTES);
            if (registeredClient.getTokenSettings().getIdTokenSignatureAlgorithm() != null) {
                jwsAlgorithm = registeredClient.getTokenSettings().getIdTokenSignatureAlgorithm();
            }


            JwtClaimsSet.Builder claimsBuilder = JwtClaimsSet.builder();
            if (StringUtils.hasText(issuer)) {
                claimsBuilder.issuer(issuer);
            }

            claimsBuilder.subject(context.getPrincipal().getName()).audience(Collections.singletonList(registeredClient.getClientId())).issuedAt(issuedAt).expiresAt(expiresAt).id(UUID.randomUUID().toString());
            SessionInformation sessionInformation;

            claimsBuilder.claim("azp", registeredClient.getClientId());
            if (AuthorizationGrantType.AUTHORIZATION_CODE.equals(context.getAuthorizationGrantType())) {
                OAuth2AuthorizationRequest authorizationRequest = context.getAuthorization().getAttribute(OAuth2AuthorizationRequest.class.getName());
                String nonce = (String)authorizationRequest.getAdditionalParameters().get("nonce");
                if (StringUtils.hasText(nonce)) {
                    claimsBuilder.claim("nonce", nonce);
                }

                sessionInformation = context.get(SessionInformation.class);
                if (sessionInformation != null) {
                    claimsBuilder.claim("sid", sessionInformation.getSessionId());
                    claimsBuilder.claim("auth_time", sessionInformation.getLastRequest());
                }
            } else if (AuthorizationGrantType.REFRESH_TOKEN.equals(context.getAuthorizationGrantType())) {
                OidcIdToken currentIdToken = context.getAuthorization().getToken(OidcIdToken.class).getToken();
                if (currentIdToken.hasClaim("sid")) {
                    claimsBuilder.claim("sid", currentIdToken.getClaim("sid"));
                }

                if (currentIdToken.hasClaim("auth_time")) {
                    claimsBuilder.claim("auth_time", currentIdToken.getClaim("auth_time"));
                }
            }

            JwsHeader.Builder jwsHeaderBuilder = JwsHeader.with(jwsAlgorithm);

            JwsHeader jwsHeader = jwsHeaderBuilder.build();
            JwtClaimsSet claims = claimsBuilder.build();
            Jwt jwt = this.jwtEncoder.encode(JwtEncoderParameters.from(jwsHeader, claims));

            return new Jwt(accessTokenGenerator.generateKey(), jwt.getIssuedAt(), jwt.getExpiresAt(), jwt.getHeaders(), jwt.getClaims());
        } else {
            return null;
        }
    }
}